#!/usr/bin/env python3
"""
展示如何使用新的数据增强功能的示例脚本
"""

import sys
import os

# 确保能找到主模块
sys.path.append('/home/dd/working/WorkpieceDetection/examples/edge_detection_robustness/external_projects/DexiNed')

from main import parse_args, main

def example_with_augmentation():
    """演示启用数据增强的训练"""
    print("=== 数据增强训练示例 ===")
    
    # 模拟命令行参数，启用数据增强
    import argparse
    
    # 创建示例参数
    example_args = [
        '--use_augmentation',           # 启用数据增强
        '--aug_flip',                   # 启用翻转
        '--aug_rotation',               # 启用旋转
        '--aug_brightness',             # 启用亮度调整
        '--aug_contrast',               # 启用对比度调整
        '--aug_noise',                  # 启用噪声
        '--aug_flip_prob', '0.5',       # 翻转概率
        '--aug_rotation_prob', '0.3',   # 旋转概率
        '--aug_brightness_prob', '0.3', # 亮度概率
        '--aug_contrast_prob', '0.3',   # 对比度概率
        '--aug_noise_prob', '0.2',      # 噪声概率
        '--aug_rotation_range', '15.0', # 旋转角度范围
        '--aug_brightness_range', '0.2',# 亮度范围
        '--aug_contrast_range', '0.2',  # 对比度范围
        '--aug_noise_std', '10.0',      # 噪声标准差
        '--epochs', '1',                # 只训练1个epoch用于测试
        '--batch_size', '2',            # 小批次用于测试
        '--is_testing', 'False'         # 训练模式
    ]
    
    # 备份原始命令行参数
    original_argv = sys.argv
    try:
        # 设置测试参数
        sys.argv = ['main.py'] + example_args
        
        # 解析参数
        args = parse_args()
        
        print("数据增强参数:")
        print(f"  启用状态: {args.use_augmentation}")
        print(f"  翻转: {args.aug_flip} (概率: {args.aug_flip_prob})")
        print(f"  旋转: {args.aug_rotation} (概率: {args.aug_rotation_prob}, 角度: ±{args.aug_rotation_range}°)")
        print(f"  亮度: {args.aug_brightness} (概率: {args.aug_brightness_prob}, 范围: ±{args.aug_brightness_range})")
        print(f"  对比度: {args.aug_contrast} (概率: {args.aug_contrast_prob}, 范围: ±{args.aug_contrast_range})")
        print(f"  噪声: {args.aug_noise} (概率: {args.aug_noise_prob}, 标准差: {args.aug_noise_std})")
        
        print("\n注意: 实际训练需要准备好数据集和模型路径")
        print("这个示例只是展示参数解析功能")
        
    finally:
        # 恢复原始命令行参数
        sys.argv = original_argv

def example_without_augmentation():
    """演示不使用数据增强的训练"""
    print("\n=== 无数据增强训练示例 ===")
    
    import argparse
    
    # 创建无增强参数
    example_args = [
        '--epochs', '1',
        '--batch_size', '2',
        '--is_testing', 'False'
        # 注意: 没有 --use_augmentation 参数
    ]
    
    original_argv = sys.argv
    try:
        sys.argv = ['main.py'] + example_args
        args = parse_args()
        
        print(f"数据增强状态: {getattr(args, 'use_augmentation', False)}")
        print("将使用传统的数据处理方式（仅随机裁剪和尺寸调整）")
        
    finally:
        sys.argv = original_argv

if __name__ == "__main__":
    print("数据增强功能使用示例")
    print("=" * 50)
    
    try:
        example_with_augmentation()
        example_without_augmentation()
        
        print("\n" + "=" * 50)
        print("✓ 参数解析测试完成!")
        print("\n使用说明:")
        print("1. 启用数据增强: python main.py --use_augmentation")
        print("2. 自定义增强参数: python main.py --use_augmentation --aug_rotation_prob 0.5 --aug_noise_std 15.0")
        print("3. 禁用特定增强: python main.py --use_augmentation --no-aug_noise")
        
    except Exception as e:
        print(f"\n✗ 测试失败: {str(e)}")
        import traceback
        traceback.print_exc()